﻿using NetMQ;
using NetMQ.Sockets;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ZeroMQCommon
{
    public class ZmqImageReceiver
    {
        private readonly ImageAssembler _assembler;
        private readonly PullSocket _socket;
        private readonly BlockingCollection<byte[]> _queue = new(10000);
        private long _receivedCount = 0;

        public ZmqImageReceiver(string bindAddress, ImageAssembler assembler)
        {
            _assembler = assembler;
            _socket = new PullSocket();
            _socket.Bind(bindAddress);
        }

        public void StartReceiving(int processorThreads = 4)
        {
            // UI stats thread
            //Task.Run(() =>
            //{
            //    while (true)
            //    {
            //        long count = Interlocked.Read(ref _receivedCount);
            //        Console.WriteLine($"ZMQ Receiver - Total Received: {count} - Memory: {GC.GetTotalMemory(false) / 1024 / 1024} MB");
            //        Thread.Sleep(1000);
            //    }
            //});

            Task.Run(() =>
            {
                while (true)
                {
                    var data = _socket.ReceiveFrameBytes();
                    Interlocked.Increment(ref _receivedCount);
                    _queue.Add(data);
                    Thread.Sleep(100);
                }
            });

           Thread.Sleep(100); // Give some time for the socket to start

            // Parallel processors
            for (int i = 0; i < processorThreads; i++)
            {
                Task.Run(() =>
                {
                    while(true)
                    {
                        try
                        {
                            if (_queue.TryTake(out var data) && data != null)
                            {
                                var meta = Encoding.UTF8.GetString(data, 0, 64).TrimEnd('\0');
                                var parts = meta.Split(':');
                                string imageId = parts[0];
                                int chunkIndex = int.Parse(parts[1]);
                                int totalChunks = int.Parse(parts[2]);
                                byte[] chunkData = new byte[data.Length - 64];
                                Buffer.BlockCopy(data, 64, chunkData, 0, chunkData.Length);
                                _assembler.AddChunk(imageId, chunkIndex, totalChunks, chunkData);
                            }

                        }
                        catch (Exception ex)
                        {
                            Console.WriteLine($"[ERROR] Failed to process chunk: {ex.Message}");
                        }
                        finally
                        {
                            Thread.Sleep(10); // Yield to avoid tight loop
                        }
                    }
                });
            }

        }
    }
}
